package es.usc.citius.hipster.examples;

import com.google.common.base.Function;
import com.google.common.base.Joiner;
import com.google.common.collect.Lists;
import com.intellij.uiDesigner.core.GridConstraints;
import com.intellij.uiDesigner.core.GridLayoutManager;
import es.usc.citius.hipster.algorithm.Hipster;
import es.usc.citius.hipster.model.CostNode;
import es.usc.citius.hipster.model.Node;
import es.usc.citius.hipster.model.Transition;
import es.usc.citius.hipster.model.function.CostFunction;
import es.usc.citius.hipster.model.function.HeuristicFunction;
import es.usc.citius.hipster.model.function.impl.StateTransitionFunction;
import es.usc.citius.hipster.model.impl.WeightedNode;
import es.usc.citius.hipster.model.problem.ProblemBuilder;
import es.usc.citius.hipster.model.problem.SearchProblem;
import es.usc.citius.hipster.util.examples.maze.Maze2D;
import es.usc.citius.hipster.util.examples.maze.Mazes;

import javax.swing.*;
import javax.swing.Timer;
import javax.swing.event.ChangeEvent;
import javax.swing.event.ChangeListener;
import java.awt.*;
import java.awt.event.ActionEvent;
import java.awt.event.ActionListener;
import java.text.DecimalFormat;
import java.util.*;
import java.util.List;

/**
 * @author Pablo Rodríguez Mier <<a href="mailto:pablo.rodriguez.mier@usc.es">pablo.rodriguez.mier@usc.es</a>>
 */
public class ASCIIMazeVisualizer {
    private JPanel mainPanel;
    private JPanel optionPanel;
    private JComboBox comboMazes;
    private JLabel loadMazeLabel;
    private JPanel algoContainerPanel;
    private JPanel textContainerPanel;
    private JTextArea mazeTextArea;
    private JComboBox comboAlgorithm;
    private JButton resetButton;
    private JButton runButton;
    private JPanel configAlgorithmPanel;
    private JSpinner refreshSpinner;
    private JPanel refreshPanel;
    private JPanel buttonPanel;
    private JLabel refreshLabel;
    private JCheckBox realtimePrintingCheckBox;
    private JPanel statusBarPanel;
    private JLabel labelInfoSteps;
    private JLabel labelSteps;
    private JLabel labelInfoCost;
    private JLabel labelCost;
    private JFrame mainFrame;
    private String lastMaze;

    {
// GUI initializer generated by IntelliJ IDEA GUI Designer
// >>> IMPORTANT!! <<<
// DO NOT EDIT OR ADD ANY CODE HERE!
        $$$setupUI$$$();
    }

    /**
     * Method generated by IntelliJ IDEA GUI Designer
     * >>> IMPORTANT!! <<<
     * DO NOT edit this method OR call it in your code!
     *
     * @noinspection ALL
     */
    private void $$$setupUI$$$() {
        mainPanel = new JPanel();
        mainPanel.setLayout(new BorderLayout(0, 0));
        optionPanel = new JPanel();
        optionPanel.setLayout(new GridLayoutManager(4, 1, new Insets(0, 0, 0, 0), -1, -1));
        optionPanel.setEnabled(true);
        mainPanel.add(optionPanel, BorderLayout.NORTH);
        loadMazeLabel = new JLabel();
        loadMazeLabel.setText("Load ASCII Maze:");
        optionPanel.add(loadMazeLabel, new GridConstraints(0, 0, 1, 1, GridConstraints.ANCHOR_WEST, GridConstraints.FILL_NONE, GridConstraints.SIZEPOLICY_FIXED, GridConstraints.SIZEPOLICY_FIXED, null, null, null, 0, false));
        comboMazes = new JComboBox();
        final DefaultComboBoxModel defaultComboBoxModel1 = new DefaultComboBoxModel();
        comboMazes.setModel(defaultComboBoxModel1);
        optionPanel.add(comboMazes, new GridConstraints(1, 0, 1, 1, GridConstraints.ANCHOR_WEST, GridConstraints.FILL_HORIZONTAL, GridConstraints.SIZEPOLICY_CAN_GROW, GridConstraints.SIZEPOLICY_FIXED, null, null, null, 0, false));
        algoContainerPanel = new JPanel();
        algoContainerPanel.setLayout(new GridLayoutManager(2, 1, new Insets(0, 0, 0, 0), -1, -1));
        optionPanel.add(algoContainerPanel, new GridConstraints(3, 0, 1, 1, GridConstraints.ANCHOR_CENTER, GridConstraints.FILL_BOTH, GridConstraints.SIZEPOLICY_CAN_SHRINK | GridConstraints.SIZEPOLICY_CAN_GROW, GridConstraints.SIZEPOLICY_CAN_SHRINK | GridConstraints.SIZEPOLICY_CAN_GROW, null, null, null, 0, false));
        comboAlgorithm = new JComboBox();
        final DefaultComboBoxModel defaultComboBoxModel2 = new DefaultComboBoxModel();
        comboAlgorithm.setModel(defaultComboBoxModel2);
        algoContainerPanel.add(comboAlgorithm, new GridConstraints(0, 0, 1, 1, GridConstraints.ANCHOR_WEST, GridConstraints.FILL_HORIZONTAL, GridConstraints.SIZEPOLICY_CAN_GROW, GridConstraints.SIZEPOLICY_FIXED, null, null, null, 0, false));
        configAlgorithmPanel = new JPanel();
        configAlgorithmPanel.setLayout(new BorderLayout(0, 0));
        algoContainerPanel.add(configAlgorithmPanel, new GridConstraints(1, 0, 1, 1, GridConstraints.ANCHOR_CENTER, GridConstraints.FILL_BOTH, GridConstraints.SIZEPOLICY_CAN_SHRINK | GridConstraints.SIZEPOLICY_CAN_GROW, GridConstraints.SIZEPOLICY_CAN_SHRINK | GridConstraints.SIZEPOLICY_CAN_GROW, null, null, null, 0, false));
        refreshPanel = new JPanel();
        refreshPanel.setLayout(new GridLayoutManager(1, 3, new Insets(0, 0, 0, 0), -1, -1));
        configAlgorithmPanel.add(refreshPanel, BorderLayout.WEST);
        refreshSpinner = new JSpinner();
        refreshPanel.add(refreshSpinner, new GridConstraints(0, 2, 1, 1, GridConstraints.ANCHOR_WEST, GridConstraints.FILL_HORIZONTAL, GridConstraints.SIZEPOLICY_WANT_GROW, GridConstraints.SIZEPOLICY_FIXED, null, new Dimension(80, -1), null, 0, false));
        refreshLabel = new JLabel();
        refreshLabel.setText("Refresh interval (ms):");
        refreshPanel.add(refreshLabel, new GridConstraints(0, 1, 1, 1, GridConstraints.ANCHOR_WEST, GridConstraints.FILL_NONE, GridConstraints.SIZEPOLICY_FIXED, GridConstraints.SIZEPOLICY_FIXED, null, null, null, 0, false));
        realtimePrintingCheckBox = new JCheckBox();
        realtimePrintingCheckBox.setSelected(true);
        realtimePrintingCheckBox.setText("Realtime printing (slow)");
        refreshPanel.add(realtimePrintingCheckBox, new GridConstraints(0, 0, 1, 1, GridConstraints.ANCHOR_WEST, GridConstraints.FILL_NONE, GridConstraints.SIZEPOLICY_CAN_SHRINK | GridConstraints.SIZEPOLICY_CAN_GROW, GridConstraints.SIZEPOLICY_FIXED, null, null, null, 0, false));
        buttonPanel = new JPanel();
        buttonPanel.setLayout(new GridLayoutManager(1, 2, new Insets(0, 0, 0, 0), -1, -1));
        configAlgorithmPanel.add(buttonPanel, BorderLayout.EAST);
        runButton = new JButton();
        runButton.setHorizontalAlignment(4);
        runButton.setText("Start");
        buttonPanel.add(runButton, new GridConstraints(0, 0, 1, 1, GridConstraints.ANCHOR_CENTER, GridConstraints.FILL_HORIZONTAL, GridConstraints.SIZEPOLICY_FIXED, GridConstraints.SIZEPOLICY_FIXED, null, null, null, 0, false));
        resetButton = new JButton();
        resetButton.setHorizontalAlignment(4);
        resetButton.setText("Reset");
        buttonPanel.add(resetButton, new GridConstraints(0, 1, 1, 1, GridConstraints.ANCHOR_CENTER, GridConstraints.FILL_HORIZONTAL, GridConstraints.SIZEPOLICY_CAN_SHRINK | GridConstraints.SIZEPOLICY_CAN_GROW, GridConstraints.SIZEPOLICY_FIXED, null, null, null, 0, false));
        final JLabel label1 = new JLabel();
        label1.setText("Select algorithm:");
        optionPanel.add(label1, new GridConstraints(2, 0, 1, 1, GridConstraints.ANCHOR_WEST, GridConstraints.FILL_NONE, GridConstraints.SIZEPOLICY_FIXED, GridConstraints.SIZEPOLICY_FIXED, null, null, null, 0, false));
        textContainerPanel = new JPanel();
        textContainerPanel.setLayout(new GridLayoutManager(1, 1, new Insets(10, 1, 1, 1), -1, -1));
        mainPanel.add(textContainerPanel, BorderLayout.CENTER);
        mazeTextArea = new JTextArea();
        mazeTextArea.setFont(new Font("Monospaced", mazeTextArea.getFont().getStyle(), 20));
        textContainerPanel.add(mazeTextArea, new GridConstraints(0, 0, 1, 1, GridConstraints.ANCHOR_CENTER, GridConstraints.FILL_BOTH, GridConstraints.SIZEPOLICY_WANT_GROW, GridConstraints.SIZEPOLICY_WANT_GROW, null, null, null, 0, false));
        statusBarPanel = new JPanel();
        statusBarPanel.setLayout(new FlowLayout(FlowLayout.LEFT, 5, 5));
        mainPanel.add(statusBarPanel, BorderLayout.SOUTH);
        statusBarPanel.setBorder(BorderFactory.createTitledBorder(BorderFactory.createLoweredBevelBorder(), null));
        labelInfoSteps = new JLabel();
        labelInfoSteps.setText("Steps:");
        statusBarPanel.add(labelInfoSteps);
        labelSteps = new JLabel();
        labelSteps.setText("0");
        statusBarPanel.add(labelSteps);
        final JSeparator separator1 = new JSeparator();
        separator1.setOrientation(1);
        statusBarPanel.add(separator1);
        labelInfoCost = new JLabel();
        labelInfoCost.setText("Distance:");
        statusBarPanel.add(labelInfoCost);
        labelCost = new JLabel();
        labelCost.setText("0");
        statusBarPanel.add(labelCost);
    }

    /**
     * @noinspection ALL
     */
    public JComponent $$$getRootComponent$$$() {
        return mainPanel;
    }

    private enum State {STOPPED, STARTED, PAUSED}

    // Global execution state
    private State state = State.STOPPED;
    // Current algorithm used
    private Iterator<? extends Node<?, Point, ?>> algorithmIterator;
    // Listener to update the text area
    private ExecutionHandler executionHandler;
    // Move all to the algorithm executor
    private Set<Point> explored = new HashSet<Point>();
    private int steps = 0;
    private Timer timer;
    private Maze2D maze;

    public static void main(String[] args) {

        try {
            UIManager.setLookAndFeel("javax.swing.plaf.nimbus.NimbusLookAndFeel");
        } catch (Exception ex) {
        }

        JFrame frame = new JFrame("Hipster Maze Shortest Path Visualizer");
        frame.setContentPane(new ASCIIMazeVisualizer(frame).mainPanel);
        frame.setDefaultCloseOperation(WindowConstants.EXIT_ON_CLOSE);
        frame.pack();
        frame.setLocationRelativeTo(null);
        frame.setVisible(true);
    }

    public ASCIIMazeVisualizer(final JFrame frame) {
        this.mainFrame = frame;
        // Use double buffer for smooth updates
        mazeTextArea.setDoubleBuffered(true);
        refreshSpinner.setValue(50);

        // Default values
        comboMazes.addItem("Maze example 1");
        comboMazes.addItem("Maze example 2");
        comboMazes.addItem("Maze example 3");
        comboMazes.addItem("Maze example 4");
        comboMazes.addItem("Maze example 5");

        comboAlgorithm.addItem("Depth First Search (DFS, non-optimal)");
        comboAlgorithm.addItem("Breadth First Search (BFS, non-optimal) ");
        comboAlgorithm.addItem("Bellman Ford");
        comboAlgorithm.addItem("Dijkstra");
        comboAlgorithm.addItem("A*");
        comboAlgorithm.addItem("IDA*");

        loadSelectedMaze();

        // Listener to process run/pause button
        runButton.addActionListener(new ActionListener() {
            @Override
            public void actionPerformed(ActionEvent e) {
                switch (state) {
                    case STARTED:
                        pause();
                        break;
                    case STOPPED:
                        start();
                        break;
                    case PAUSED:
                        continueExecution();
                        break;
                }
            }
        });

        resetButton.addActionListener(new ActionListener() {
            @Override
            public void actionPerformed(ActionEvent e) {
                stop();
                loadSelectedMaze();
            }
        });

        comboMazes.addActionListener(new ActionListener() {
            @Override
            public void actionPerformed(ActionEvent e) {
                loadSelectedMaze();
                // Resize to adapt the window
                frame.pack();
            }
        });


        refreshSpinner.addChangeListener(new ChangeListener() {
            @Override
            public void stateChanged(ChangeEvent e) {
                timer.setDelay((Integer) refreshSpinner.getValue());
            }
        });

        // Initialize the thread / listener for the execution
        executionHandler = new ExecutionHandler();
        timer = new Timer((Integer) refreshSpinner.getValue(), executionHandler);
        new Thread(executionHandler).start();
        timer.start();
    }

    private void continueExecution() {
        if (state.equals(State.PAUSED)) {
            runButton.setText("Pause");
            state = State.STARTED;
        }
    }

    private synchronized Node<?, Point, ?> executeSearchStep() {
        Node<?, Point, ?> node = null;
        if (algorithmIterator.hasNext()) {
            node = algorithmIterator.next();
            steps++;
            explored.add(node.state());
            if (realtimePrintingCheckBox.isSelected()) updateVisualizer(node, maze, explored);
            if (node.state().equals(maze.getGoalLoc())) {
                updateVisualizer(node, maze, explored);
                stop();
            }
        } else {
            updateVisualizer(node, maze, explored);
            stop();
        }
        return node;
    }

    private void start() {
        // Create a new maze and run the selected algorithm
        steps = 0;
        try {
            maze = new Maze2D(mazeTextArea.getText().split("\\r?\\n"));
        } catch (Exception ex) {
            JOptionPane.showMessageDialog(mainFrame, ex.getMessage() + ". Try to reset the map.", "Maze parse exception", JOptionPane.ERROR_MESSAGE);
            return;
        }
        // Create a new algorithm
        algorithmIterator = createAlgorithm(maze);
        // Reset explored tiles
        this.explored = new HashSet<Point>();
        runButton.setText("Pause");
        state = State.STARTED;
    }

    private void stop() {
        runButton.setText("Start");
        state = State.STOPPED;
    }

    private void pause() {
        runButton.setText("Resume");
        state = State.PAUSED;
    }

    private synchronized void updateVisualizer(final Node<?, Point, ?> node, Maze2D maze, Collection<Point> explored) {
        if (node != null && maze != null) {
            List<Point> statePath = Lists.transform(node.path(), new Function<Node<?, Point, ?>, Point>() {
                @Override
                public Point apply(Node<?, Point, ?> pointNode) {
                    return pointNode.state();
                }
            });

            String mazeStr = getMazeStringSolution(maze, explored, statePath);
            mazeTextArea.setText(mazeStr);
            // Update the status bar
            SwingUtilities.invokeLater(new Runnable() {
                @Override
                public void run() {
                    labelSteps.setText(Integer.toString(steps));
                    if (node instanceof CostNode) {
                        CostNode n = (CostNode) node;
                        labelCost.setText(new DecimalFormat("#.00").format(n.getCost()));
                    }
                }
            });
        }
    }


    private class ExecutionHandler implements ActionListener, Runnable {

        private ExecutionHandler() {
        }

        @Override
        public void actionPerformed(ActionEvent e) {
            if (realtimePrintingCheckBox.isSelected() && state.equals(State.STARTED)) {
                executeSearchStep();
            }
        }

        @Override
        public void run() {
            while (!Thread.currentThread().isInterrupted()) {
                try {
                    if (!realtimePrintingCheckBox.isSelected() && state.equals(State.STARTED)) {
                        executeSearchStep();
                    } else {
                        Thread.sleep(100);
                        Thread.yield();
                    }
                } catch (InterruptedException ex) {
                    Thread.currentThread().interrupt();
                }
            }
        }
    }

    private Iterator<? extends Node<?, Point, ?>> createAlgorithm(Maze2D maze) {
        final Iterator<? extends Node<?, Point, ?>> iterator;
        switch (comboAlgorithm.getSelectedIndex()) {

            case 0:
                iterator = Hipster.createDepthFirstSearch(buildProblem(maze, false)).iterator();
                break;
            case 1:
                iterator = Hipster.createBreadthFirstSearch(buildProblem(maze, false)).iterator();
                break;
            case 2:
                iterator = Hipster.createBellmanFord(buildProblem(maze, false)).iterator();
                break;
            case 3:
                iterator = Hipster.createDijkstra(buildProblem(maze, false)).iterator();
                break;
            case 4:
                iterator = Hipster.createAStar(buildProblem(maze, true)).iterator();
                break;
            case 5:
                iterator = Hipster.createIDAStar(buildProblem(maze, true)).iterator();
                break;
            default:
                throw new IllegalStateException("Invalid algorithm");
        }
        return iterator;
    }

    private void loadSelectedMaze() {
        switch (comboMazes.getSelectedIndex()) {
            case 0:
                mazeTextArea.setText(Joiner.on('\n').join(Mazes.exampleMaze1));
                break;
            case 1:
                mazeTextArea.setText(Joiner.on('\n').join(Mazes.testMaze4));
                break;
            case 2:
                mazeTextArea.setText(Joiner.on('\n').join(Mazes.testMaze3));
                break;
            case 3:
                mazeTextArea.setText(Joiner.on('\n').join(Mazes.testMaze2));
                break;
            case 4:
                mazeTextArea.setText(Joiner.on('\n').join(Mazes.testMaze5));
                break;
        }
    }

    private String getMazeStringSolution(Maze2D maze, Collection<Point> explored, Collection<Point> path) {
        List<Map<Point, Character>> replacements = new ArrayList<Map<Point, Character>>();
        Map<Point, Character> replacement = new HashMap<Point, Character>();
        for (Point p : explored) {
            replacement.put(p, '.');
        }
        replacements.add(replacement);
        replacement = new HashMap<Point, Character>();
        for (Point p : path) {
            replacement.put(p, '*');
        }
        replacements.add(replacement);
        return maze.getReplacedMazeString(replacements);
    }

    private SearchProblem<Void, Point, WeightedNode<Void, Point, Double>> buildProblem(final Maze2D maze, final boolean heuristic) {
        return ProblemBuilder.create()
                .initialState(maze.getInitialLoc())
                .defineProblemWithoutActions()
                .useTransitionFunction(new StateTransitionFunction<Point>() {
                    @Override
                    public Iterable<Point> successorsOf(Point state) {
                        return maze.validLocationsFrom(state);
                    }
                })
                .useCostFunction(new CostFunction<Void, Point, Double>() {
                    @Override
                    public Double evaluate(Transition<Void, Point> transition) {
                        Point source = transition.getFromState();
                        Point destination = transition.getState();
                        double distance = source.distance(destination);
                        double roundedDistance = (double) Math.round(distance * 1e5) / 1e5;
                        return roundedDistance;
                    }
                })
                .useHeuristicFunction(new HeuristicFunction<Point, Double>() {
                    @Override
                    public Double estimate(Point state) {
                        if (heuristic) {
                            double distance = state.distance(maze.getGoalLoc());
                            double roundedDistance = (double) Math.round(distance * 1e5) / 1e5;
                            return roundedDistance;
                        }
                        return 0d;
                    }
                })
                .build();
    }
}
